#ifndef __threadPool_H__
#define __threadPool_H__

#include <iostream>
#include <thread>
#include <mutex>
#include <queue>
#include <condition_variable>
#include <functional>
#include <stdlib.h>
#include <assert.h>
//线程安全队列
template<typename T>
class CMyQueue
{
public:
    void push(const T& item)
    {
        {
            std::scoped_lock lock(m_mutex);
            m_queue.push(item);
        }
        m_cond.notify_one();
    }

    void push(T&& item)
    {
        {
            std::scoped_lock lock(m_mutex);
            m_queue.push(std::move(item));
        }
        m_cond.notify_one();
    }

    bool pop(T &item)
    {
        std::unique_lock lock(m_mutex);
        m_cond.wait(lock, [&](){
            return !m_queue.empty() || m_stop;
        });

        if(m_queue.empty())
        {
            return false;
        }

        item = std::move(m_queue.front());
        m_queue.pop();
        return true;
    }

    std::size_t size() const
    {
        std::scoped_lock lock(m_mutex);
        return m_queue.size();
    }

    bool empty() const
    {
        std::scoped_lock lock(m_mutex);
        return m_queue.empty();
    }

    void stop()
    {
        {
            std::scoped_lock lock(m_mutex);
            m_stop = true;
        }
        m_cond.notify_all();
    }
private:
    std::condition_variable m_cond;
    std::queue<T> m_queue;
    mutable std::mutex m_mutex;
    bool m_stop = false;
};


class CThreadPoll
{
public:
    explicit CThreadPoll(size_t threads = std::thread::hardware_concurrency()):
    m_queue(threads),
    m_threadNum(threads)
    {
        auto worker = ([this](size_t id){
            while(1)
            {
                std::function<void()> task;
                if(!m_queue[id].pop(task))
                {
                    break;
                }

                if(task)
                {
                    task();
                }
            }
         });

        m_threadWorker.reserve(threads);
        for(int i = 0; i < threads; ++i)
        {
            m_threadWorker.emplace_back(worker, i);
        }
        
    }

    int enqueue(std::function<void()> item, size_t id = -1)
    {
        if(item == nullptr)
        {
            return -1;      
        }

        if(id == -1)
        {
            int n = rand() % m_threadNum;
            m_queue[n].push(std::move(item));
        }
        else
        {
            assert(id < m_threadNum);
            m_queue[id].push(std::move(item));
        }
        
        return 0;

    }
    
    size_t getThreadNum()
    {
        return m_threadNum;
    }

    ~CThreadPoll()
    {
        for(auto& q: m_queue)
        {
            q.stop();
        }

        for(auto& w: m_threadWorker)
        {
            w.join();
        }
    }
private:
    std::vector<CMyQueue<std::function<void()>>> m_queue;
    std::vector<std::thread> m_threadWorker;
    size_t m_threadNum;
};

#endif
